import fitz  # PyMuPDF
import io
import os
from PIL import Image
import sys
sys.path.append('/home/FAST_DATA_MIRROR/Langchain-Chatchat-master')
from pdf_tools.duguang import Duguang
class ReadImInPdf:
    def __init__(self) -> None:
        self.duguang = Duguang()
    def pdf_image_extract(self, path):
        doc = fitz.open(path)
        
        for i in range(len(doc)):
            for j, img in enumerate(doc.get_page_images(i)):
                xref = img[0]
                base = img[1]
                img_data = doc.extract_image(xref)
                img_bytes = img_data['image']
                image = Image.open(io.BytesIO(img_bytes))
                colors = image.getcolors()
                if not colors:
                    continue
                    # image.save(f"/home/FAST_DATA_MIRROR/Langchain-Chatchat-master/pdf_tools/output/none/{i}-{j}.png")
                # 判断是否为纯色
                elif len(colors) == 1:
                    image.save(f"/home/FAST_DATA_MIRROR/Langchain-Chatchat-master/pdf_tools/output/none/{i}-{j}.png")
                else:
                    image_cv2 = self.duguang.infer(image)
                    Image.fromarray(image_cv2).save(f"/home/FAST_DATA_MIRROR/Langchain-Chatchat-master/pdf_tools/output/text/{i}-{j}.png")
if __name__ == '__main__':
    pdf_image = ReadImInPdf()
    pdf_image.pdf_image_extract('/home/FAST_DATA_MIRROR/Langchain-Chatchat-master/云鹰平台服务系统接口协议V1.8.4.pdf')